package com.zlz.m3u8download.controller;

import com.zlz.m3u8download.bean.ResultSet;
import com.zlz.m3u8download.bean.RunTimeTask;
import com.zlz.m3u8download.task.CommonCache;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 不同的userId会创建不同的实例
 * 相同的userId在@OnMessage每次请求接收到消息时，使用的同一个实例接收
 */
@ServerEndpoint(value = "/websocket/{userId}/{taskId}")
@Component
@Slf4j
public class TaskWebSocket {
    /**
     * 存放所有在线的客户端
     */
    public static final ConcurrentHashMap<String, Session> webSocketMap = new ConcurrentHashMap<>();

    /**
     * 关联
     */
    public static final ConcurrentHashMap<String, String> userTaskMap = new ConcurrentHashMap<>();

    /**
     * 关联
     */
    public static final ConcurrentHashMap<String, Set<String>> taskUserMap = new ConcurrentHashMap<>();

    public static void sendTaskMessage(String userId, String taskId) {
        // 任务等待
        if (CommonCache.priorityTask.contains(Long.valueOf(taskId))) {
            TaskWebSocket.send(userId, new ResultSet<>("success", 0).toJSONString());
            return;
        }
        if (CommonCache.waitingTasks.contains(Long.valueOf(taskId))) {
            TaskWebSocket.send(userId, new ResultSet<>("success", 0).toJSONString());
            return;
        }
        RunTimeTask runTimeTask = CommonCache.runningTasks.get(Long.valueOf(taskId));
        // 任务完成
        if (runTimeTask == null) {
            TaskWebSocket.send(userId, new ResultSet<>("success", 2).toJSONString());
            return;
        }
        // 任务进行中
        runTimeTask.setMinCode(runTimeTask.getWaitingTaskCode().iterator().next());
        TaskWebSocket.send(userId, new ResultSet<>("success", 1, runTimeTask).toJSONString());
    }

    //新增一个方法用于主动向客户端发送消息
    public static void send(String userId, String message) {
        Session session = webSocketMap.get(userId);
        if (session != null) {
            session.getAsyncRemote().sendText(message);
        }
    }


    //前端请求时一个websocket时
    @OnOpen
    public void onOpen(Session session,
                       @PathParam("taskId") String taskId,
                       @PathParam("userId") String userId) {
        Set<String> list = taskUserMap.computeIfAbsent(taskId, (a) -> new HashSet<>());
        list.add(userId);
        webSocketMap.put(userId, session);
        userTaskMap.put(userId, taskId);
        log.info("【websocket消息】有新的连接,userId={},taskId={}", userId, taskId);
        sendTaskMessage(userId, taskId);
    }

    //前端关闭时一个websocket时
    @OnClose
    public void onClose(@PathParam("userId") String userId, @PathParam("taskId") String taskId) {
        webSocketMap.remove(userId);
        userTaskMap.remove(userId);
        Set<String> userIds = taskUserMap.get(taskId);
        if (userIds != null) {
            userIds.remove(userId);
        }
        log.info("【websocket消息】连接断开:{}", userId);
    }

    @OnError
    public void onError(Session session, Throwable error) {
        log.info("【websocket消息】WebSocket发生错误，错误信息为：{}", error.getMessage());
        log.error("错误", error);
    }

    //前端向后端发送消息
    @OnMessage
    public void onMessage(Session session, String message,
                          @PathParam("taskId") String taskId,
                          @PathParam("userId") String userId) {
        sendTaskMessage(userId, taskId);
    }
}
 